import os
import pandas as pd
import glob
import plotly.graph_objects as go
from plotly.offline import plot

class FileUtils:
    @staticmethod
    def save_to_excel(data, folder_path, filename=None):
        """
        将数据保存到Excel文件
        
        Args:
            data: DataFrame数据
            folder_path: 保存文件夹路径
            filename: 文件名，默认为当前日期
        
        Returns:
            保存的文件路径
        """
        # 确保文件夹存在
        os.makedirs(folder_path, exist_ok=True)
        
        # 如果没有指定文件名，使用当前日期
        if filename is None:
            filename = pd.Timestamp.now().strftime('%Y%m%d')
        
        # 构建完整的文件路径
        file_path = f"{folder_path}/{filename}.xlsx"
        
        # 保存数据到Excel
        data.to_excel(file_path, index=False)
        
        return file_path
    
    @staticmethod
    def generate_heatmap_html(data_folder, output_path, colorscale='Viridis'):
        """
        生成热力图HTML文件
        
        Args:
            data_folder: 数据文件夹路径
            output_path: 输出HTML文件路径
            colorscale: 颜色方案，可以是字符串或自定义颜色列表
        
        Returns:
            输出的HTML文件路径
        """
        # 获取所有xlsx文件
        xlsx_files = glob.glob(os.path.join(data_folder, "*.xlsx"))
        xlsx_files.sort()  # 按文件名排序
        
        # 只保留最近30天的文件
        if len(xlsx_files) > 30:
            xlsx_files = xlsx_files[-30:]
        
        if not xlsx_files:
            print("没有找到xlsx文件")
            return None
        
        # 读取所有数据
        all_data = []
        dates = []
        
        for file in xlsx_files:
            try:
                # 从文件名提取日期
                date_str = os.path.basename(file).split('.')[0]
                date = pd.to_datetime(date_str, format='%Y%m%d')
                dates.append(date.strftime('%Y-%m-%d'))
                
                # 读取数据
                df = pd.read_excel(file)
                all_data.append(df)
            except Exception as e:
                print(f"读取文件 {file} 时出错: {e}")
        
        if not all_data:
            print("没有成功读取任何数据")
            return None
        
        # 合并所有板块
        all_sectors = set()
        for df in all_data:
            all_sectors.update(df['板块'].tolist())
        
        all_sectors = sorted(list(all_sectors))
        
        # 创建热力图数据
        z_data = []
        sector_latest_scores = {}  # 用于存储最新一天的得分
        
        for sector in all_sectors:
            row = []
            for i, df in enumerate(all_data):
                sector_data = df[df['板块'] == sector]
                if not sector_data.empty:
                    score = sector_data['热力得分'].values[0]
                    row.append(score)
                    # 记录最新一天的得分
                    if i == len(all_data) - 1:
                        sector_latest_scores[sector] = score
                else:
                    row.append(None)  # 缺失数据
            z_data.append(row)
        
        # 根据最新一天的得分对板块进行排序
        sorted_sectors_with_scores = sorted(sector_latest_scores.items(), key=lambda x: x[1], reverse=False)  # 改为 False，使高分在上
        sorted_sectors = [item[0] for item in sorted_sectors_with_scores]
        
        # 重新排序数据
        sorted_z_data = []
        for sector in sorted_sectors:
            idx = all_sectors.index(sector)
            sorted_z_data.append(z_data[idx])
        
        # 创建热力图
        fig = go.Figure(data=go.Heatmap(
            z=sorted_z_data,
            x=dates,
            y=sorted_sectors,  # 使用排序后的板块
            colorscale=colorscale,
            colorbar=dict(title='热力得分'),
            hoverongaps=False
        ))
        
        # 添加文本注释
        annotations = []
        for i, sector in enumerate(sorted_sectors):  # 使用排序后的板块
            for j, date in enumerate(dates):
                idx = all_sectors.index(sector)
                if z_data[idx][j] is not None:  # 确保有数据
                    score = z_data[idx][j]
                    # 计算文字颜色 - 根据热力值决定
                    if score < 60 or score > 80:
                        text_color = 'white'
                    else:
                        text_color = 'black'
                    
                    annotations.append(dict(
                        x=date,
                        y=sector,
                        text=f"{score:.1f}",
                        showarrow=False,
                        font=dict(
                            color=text_color,
                            size=10
                        ),
                        align='center'
                    ))
        # 设置布局
        fig.update_layout(
            title='板块热力得分变化趋势',
            xaxis=dict(title='日期'),
            yaxis=dict(title='板块'),
            height=max(800, len(all_sectors) * 20),  # 最小高度800像素，每个板块高度20像素
            width=max(1200, len(dates) * 44),  # 最小宽度1200像素，每个日期宽度44像素
            annotations=annotations  # 添加注释
        )
        
        # 保存为HTML文件
        plot(fig, filename=output_path, auto_open=False)
        
        return output_path